#!/usr/bin/python

from numpy import *
import exr, sys, os
from optparse import OptionParser
from kernel import Kernel
from misc import mk_frame, find_index
from gui import GlGui

import pycuda.autoinit
import pycuda.driver as cuda


if __name__ == '__main__':
    main()


def main():
    # read arguments and options
    if len(sys.argv) <= 1:
        print "not enough arguments, try shadie --help"
        sys.exit()

    parser = OptionParser()
    parser.add_option("-o", "--out", dest="output_file", default=None,
                      help="write output to EXRFILE", metavar="EXRFILE")
    parser.add_option("-c", "--camera-position", dest="camera_pos", default="1,1,1",
                      help="set camera position to x,y,z", metavar="x,y,z")
    parser.add_option("-u", "--camera-up", dest="camera_up", default="0,0,1",
                      help="set camera up-vector to x,y,z", metavar="x,y,z")
    parser.add_option("-z", "--camera-zoom", dest="camera_scale", default="1",
                      help="set camera scaling to x", metavar="x")
    parser.add_option("-s", "--image-size", dest="image_size", default="512",
                      help="set image size to N x N", metavar="N")
    parser.add_option("-t", "--time", dest="time", default="0",
                      help="set time to T", metavar="T")
    parser.add_option("-R", "--recording-prefix", dest="recording_prefix",
            default="", help="prefix for recorded frames (use F8)",
            metavar="PREFIX")
    parser.add_option("-r", "--param-resolution", dest="param_resolution", default="20",
                      help="set time to T", metavar="T")
    parser.add_option("-v", "--verbose", dest="verbose", default=False, help="print additional info")
    parser.add_option("-p", "--progressive", dest="progressive", action='store_true', default=False)

    (options, args) = parser.parse_args()

    camera_pos = array([float(s) for s in options.camera_pos.split(",")], float32)
    camera_up = array([float(s) for s in options.camera_up.split(",")], float32)
    assert camera_pos.shape == (3,)
    assert camera_up.shape == (3,)
    camera_scale = float32(options.camera_scale)


    input_file = args[0]
    delete_tmp_file = False

    if input_file.endswith('.py'):
        # invoke compiler
        shader_file = '/tmp/' + input_file[:-3] + '.cu'
        delete_tmp_file = True
        result = os.system('py2cu ' + input_file + ' ' + shader_file)
        if result != 0:
            sys.exit()
    elif input_file.endswith('.cu'):
        shader_file = input_file
    else:
        print 'wrong input file extension'
        sys.exit()

    # setup kernel
    kernel = Kernel(shader_file, options)

    if delete_tmp_file:
        os.remove(shader_file)

    # read command-line parameter overrides
    for key_val in args[1:]:
        key, val = key_val.split('=')
        print key, val
        idx = find_index(lambda x: x.name == key, kernel.params)
        if idx != None:
            kernel.params[idx].value = float(val)
        else:
            idx = find_index(lambda x: x.name == key, kernel.datasets)
            assert idx != None, "no such parameter or dataset: " + key
            kernel.datasets[idx].filename = val
        
    # load CUDA module and datasets
    kernel.load()

    img_width = int(options.image_size)
    img_height = int(options.image_size)

    if options.output_file != None:
        # run the kernel once and write to file
        pstruct = kernel.create_params_struct()
        img = zeros((img_height, img_width, 3), dtype = float32, order = "C")
        
        frame = mk_frame(camera_pos, r_[0, 0, 0], camera_up)
        frame = array(frame, dtype=float32)
        
        kernel.run(float32(options.time), int32(img_width), int32(img_height), cuda.In(frame),
            camera_scale, float32(-1), cuda.In(pstruct), cuda.Out(img), intp(0),
            block = (8, 8, 1), grid = (img_width/8 + 1, img_height/8 + 1))
            
        # flip
        img = img[::-1, :, :]
            
        exr.save(img, options.output_file)
        
    else:
        # start glut GUI
        gl_gui = GlGui(kernel, options)
        gl_gui.start()
        

